import os
from dotenv import load_dotenv
import yaml
from litellm.router import Router
import logging

class ModelRouter:
    def __init__(self, config_path=None, env_path=None):
        """Initialize model router with configuration"""
        # 获取项目根目录
        project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
        
        # 设置默认配置文件路径
        if config_path is None:
            config_path = os.path.join(project_root, "config", "config.yaml")
        elif not os.path.isabs(config_path):
            config_path = os.path.join(project_root, config_path)
            
        if env_path is None:
            env_path = os.path.join(project_root, "config", ".env")
        elif not os.path.isabs(env_path):
            env_path = os.path.join(project_root, env_path)
        
        # 验证文件是否存在
        if not os.path.exists(config_path):
            raise FileNotFoundError(f"Configuration file not found: {config_path}")
        if not os.path.exists(env_path):
            raise FileNotFoundError(f"Environment file not found: {env_path}")
            
        load_dotenv(dotenv_path=env_path)
        with open(config_path, 'r', encoding='utf-8') as file:
            self.config = yaml.safe_load(file)
        self.router = self._initialize_router()
    
    def _initialize_router(self):
        """Create LiteLLM router instance with load balancing"""
        model_list = []
        for model in self.config['models']:
            # Use the full model name directly from config
            full_model_name = model['name']
            api_key = os.getenv(model['api_key_env'])
            
            # Validate API key exists
            if not api_key:
                raise ValueError(f"API key not found for environment variable: {model['api_key_env']}")
            
            model_config = {
                "model_name": full_model_name,
                "litellm_params": {
                    "model": full_model_name,
                    "api_key": api_key,
                    **model.get('params', {})
                }
            }
            model_list.append(model_config)
        return Router(model_list=model_list)
    
    def completion(self, **kwargs):
        """Route completion request with load balancing"""
        return self.router.completion(**kwargs)
    
    async def acompletion(self, **kwargs):
        """Async version of completion"""
        return await self.router.acompletion(**kwargs)